import os
import sys
import numpy as np

def gen_data(shape, dtype, rand_type, low, high):
    if rand_type == "randint":
        rand_data = np.random.randint(low, high, size=shape)
    else:
        rand_data = np.random.uniform(low, high, size=shape)
    data = np.array(rand_data, dtype=dtype)

    return data

def gen_gathernd_data(shape, input_one_path, input_two_path):
    bs = shape[0]
    size = shape[1]

    data = gen_data(shape, np.float32, "randfloat", 0, size)
    index = gen_data([10*bs, 2], np.int32, "randint", 0, size)

    data.astype(np.float32).tofile(input_one_path)
    index.astype(np.int32).tofile(input_two_path)

def main():
    work_dir = os.getcwd()
    data_dir = os.path.join(work_dir, "data")

    ## Generate GatherND op inputs
    bs = 100
    size = 21128
    input_one_path = os.path.join(data_dir, "inputs", "gathernd_params_bs%d.bin" % bs)
    input_two_path = os.path.join(data_dir, "inputs", "gathernd_indices_bs%d.bin" % bs)
    gen_gathernd_data([bs, size], input_one_path, input_two_path)
    print("generate gather_nd input data ok, bs:%d" % bs)

if __name__ == "__main__":
    main()